from util import ClickhouseUtil, JsonUtil, PostSender
import argparse


class TextGen:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.parser.add_argument("-i", "--index", type=int, help="index of the task", default=-1)
        self.parser.add_argument("-m", "--model_dir", type=str, help="model directory",
                                 default="/home/user1/models/llama-beta2.0/")
        self.parser.add_argument("-o", "--output", type=str, help="output file directory",
                                 default="/home/user1/jdigs_benchmark/data/")
        args = self.parser.parse_args()

        self.model_dir = args.model_dir
        self.current_index = args.index
        # 模型地址需要确保是以/结尾
        self.model_name = self.model_dir.split('/')[-2]
        self.output_file = args.output
        self.llm_url = "http://localhost:8000/v1/completions"
        self.clickhouse_util = ClickhouseUtil()

        self.PROMPT_TEMPLATE = ("你将完成生成中国司法裁判文书中a判决书的任务，"
                                "该类文书分为以下几段：b、引用法条段、判决段。"
                                "接下来将提供c段，请你生成d段，不要生成其他段落。c段如下：")
        self.key_words = [
            ['民事一审', '证据段、事实段、裁判分析段', '证据', '事实'],
            ['民事一审', '证据段、事实段、裁判分析段', '事实', '裁判分析'],
            ['民事一审', '证据段、事实段、裁判分析段', '事实', '引用法条'],
            ['民事一审', '证据段、事实段、裁判分析段', '事实', '判决'],

            ['民事二审', '证据事实段、裁判分析段', '证据事实', '裁判分析'],
            ['民事二审', '证据事实段、裁判分析段', '证据事实', '引用法条'],
            ['民事二审', '证据事实段、裁判分析段', '证据事实', '判决'],

            ['刑事一审', '证据段、分析段', '证据', '分析'],
            ['刑事一审', '证据段、分析段', '分析', '引用法条'],
            ['刑事一审', '证据段、分析段', '分析', '判决']
        ]
        self.QUERYs = [
            'SELECT t.evidence,t.fact FROM judgement.civil1 t',
            'SELECT t.fact,t.analysis FROM judgement.civil1 t',
            'SELECT t.fact,t.law FROM judgement.civil1 t',
            'SELECT t.fact,t.judgement FROM judgement.civil1 t',

            'SELECT t.evidence,t.analysis FROM judgement.civil2 t',
            'SELECT t.evidence,t.law FROM judgement.civil2 t',
            'SELECT t.evidence,t.judgement FROM judgement.civil2 t',

            'SELECT t.evidence,t.analysis FROM judgement.criminal1 t',
            'SELECT t.analysis,t.law FROM judgement.criminal1 t',
            'SELECT t.analysis,t.judgement FROM judgement.criminal1 t',
        ]

        self.task_name = [
            "civil1_evidence2fact", "civil1_fact2analysis",
            "civil1_fact2law", "civil1_fact2judgement",
            "civil2_evidence2analysis", "civil2_evidence2law",
            "civil2_evidence2judgement", "criminal1_evidence2analysis",
            "criminal1_analysis2law", "criminal1_analysis2judgement"
        ]
        self.column_names = ["prompt", "original", "generated"]
        self.database = "benchmark"

    def get_origin_text(self):
        key_word = self.key_words[self.current_index]
        prompt = self.PROMPT_TEMPLATE
        prompt = prompt.replace('a', key_word[0])
        prompt = prompt.replace('b', key_word[1])
        prompt = prompt.replace('c', key_word[2])
        prompt = prompt.replace('d', key_word[3])

        result = self.clickhouse_util.get_data_from_clickhouse(query=self.QUERYs[self.current_index])
        return result, prompt

    def save_origin_and_generated_text(self, result, prompt_prefix):
        sender = PostSender(self.llm_url, self.model_dir)
        json_util = JsonUtil(self.get_json_file_name())
        clickhouse_data = []
        cnt = 0

        for row in result:
            if cnt == 500:
                break
            need = row[0]
            original = row[1]
            prompt = prompt_prefix + need

            # 发送请求，获取生成的文本
            generated = sender.send(prompt)

            data = {
                "model_name": self.model_name,
                "prompt": prompt,
                "original": original,
                "generated": generated
            }

            # 如果generated为空，则不写入文件
            if generated is None or len(generated) == 0:
                print("生成结果为空")
                continue
            else:
                json_util.json_write(data)
                clickhouse_data.append([prompt, original, generated])
                cnt += 1
                print(f"已生成{cnt}条数据")

        self.clickhouse_util.write_data(self.get_clickhouse_table_name(), clickhouse_data, self.column_names, self.database)

    def get_clickhouse_table_name(self):
        return self.task_name[self.current_index] + "_" + self.model_name

    def get_json_file_name(self):
        return self.output_file + self.model_name + "/" + self.task_name[self.current_index] + ".jsonl"


if __name__ == '__main__':
    text_gen = TextGen()
    if text_gen.current_index == -1:
        for i in range(0, len(text_gen.task_name)):
            text_gen.current_index = i
            result, prompt_prefix = text_gen.get_origin_text()
            text_gen.save_origin_and_generated_text(result, prompt_prefix)
    else:
        result, prompt_prefix = text_gen.get_origin_text()
        text_gen.save_origin_and_generated_text(result, prompt_prefix)

